package com.bigdata.spark.ml
import org.apache.spark.sql.SparkSession
import org.apache.spark.ml.linalg.{Vector, Vectors}
import org.apache.spark.ml.Pipeline
import org.apache.spark.ml.feature.{IndexToString, StringIndexer, VectorIndexer}

import org.apache.spark.ml.classification.DecisionTreeClassificationModel
import org.apache.spark.ml.classification.DecisionTreeClassifier
import org.apache.spark.ml.evaluation.MulticlassClassificationEvaluator

/**
 * @author Gerry chan
 * @version 1.0
 * 决策树
 * 课程地址：https://www.bilibili.com/video/BV1Yt411A7o6?p=8
 *
 */
object DecisionTreeIris {

  case class Iris(features:org.apache.spark.ml.linalg.Vector,label:String)

  def main(args: Array[String]): Unit = {
    val spark = SparkSession.builder()
      .master("local[*]")
      .appName("DecisionTreeIris")
      .getOrCreate()

    import spark.implicits._

    //读取数据，简要分析
    val data = spark.sparkContext.textFile("datas/iris.data.txt")
      .map(_.split(","))
      .map(p =>{
        Iris(Vectors.dense(p(0).toDouble, p(1).toDouble,p(2).toDouble,
          p(3).toDouble), p(4).toString())
      }).toDF()

    data.createOrReplaceTempView("iris")
    val df = spark.sql("select * from iris")
    //查看下数据
    df.map(t => t(1) + ":" + t(0)).collect().foreach(println)

    //进一步处理特征和标签，以及数据分组
    //分别获取标签列和特征列，进行索引，并进行了重命名
    val labelIndexer = new StringIndexer()
      .setInputCol("label")
      .setOutputCol("indexedLabel")
      .fit(df)

    val featureIndexer = new VectorIndexer()
      .setInputCol("features")
      .setOutputCol("indexedFeatures")
      .setMaxCategories(4)
      .fit(df)

    //这里我们设置一个labelConverter,目的是把预测的类型重新转换成字符型的
    val labelConverter = new IndexToString()
      .setInputCol("prediction")
      .setOutputCol("predictedLabel")
      .setLabels(labelIndexer.labels)

    //接下来，我们把数据集随意分成训练集和测试集，其中训练集占70%
    val Array(trainingData, testData) = data.randomSplit(Array(0.7,0.3))

    //构建决策树分类模型
    //训练决策树模型，这里我们可以通过setter的方法来设置决策树的参数，
    //也可以用 ParamMap来设置(具体的可以查看spark mllib官网)。
    //具体的可以设置的参数可以通过explainParams()来获取
    val dtClassifier = new DecisionTreeClassifier()
      .setLabelCol("indexedLabel")
      .setFeaturesCol("indexedFeatures")

    //定义评估器， 在pipeline中设置
    val pipelinedClassifier = new Pipeline()
      .setStages(Array(labelIndexer, featureIndexer, dtClassifier, labelConverter))

    //训练决策树模型
    val modelClassfier = pipelinedClassifier.fit(trainingData)

    //进行预测
    val predictionsClassfier = modelClassfier.transform(testData)

    //查看部分的预测结果
    predictionsClassfier.select("predictedLabel", "label", "features").show(10)

    //评估决策树分类模型
    val evaluatorClassfier = new MulticlassClassificationEvaluator()
      .setLabelCol("indexedLabel")
      .setPredictionCol("prediction")
      .setMetricName("accuracy")

    val accuracy  = evaluatorClassfier.evaluate(predictionsClassfier)

    println("Test Error= " + (1.0 - accuracy))

    //获取整个决策树模型
    val treeModelClassifier = modelClassfier.stages(2).asInstanceOf[DecisionTreeClassificationModel]

    println("Learned classification tree model:\n" +
      treeModelClassifier.toDebugString
    )
  }

}
